#include <filesystem>
#include <iostream>

#include "filter.h"

const std::string xml_dir_path = "dataset/image_annotation/";
const std::string img_dir_path = "dataset/image/";
const std::string output_path = "train/";

int count[10];

void processImg(const std::string &name) {
  std::string xml_path = xml_dir_path + name + ".xml";
  std::string img_path = img_dir_path + name + ".jpg";
  std::vector<ArmorAnnotation> annotation = parseXmlFile(xml_path);
  cv::Mat img = cv::imread(img_path);
  if (img.empty()) {
    std::cerr << "未找到：" << img_path << std::endl;
    return;
  }

  for (auto &ann : annotation) {
    if (ann.bbox.x < 0 || ann.bbox.y < 0 ||
        ann.bbox.x + ann.bbox.width > img.cols ||
        ann.bbox.y + ann.bbox.height > img.rows) {
      continue; // 忽略无效 BBox
    }

    cv::Mat roi_img = img(ann.bbox);
    cv::Mat wraped_armor;
    bool success = findAndWarpArmor(roi_img, wraped_armor, ann.color);
    if (success) {
      ++count[ann.class_id];
      std::string output_name = output_path + std::to_string(ann.class_id) +
                                "/" + std::to_string(count[ann.class_id]) +
                                ".jpg";
      cv::imwrite(output_name, wraped_armor);
    }
  }
}

namespace fs = std::filesystem;

std::vector<std::string> getXmlFilePaths(const std::string &directory_path) {
  std::vector<std::string> xml_paths;

  // 检查目录是否存在
  if (!fs::exists(directory_path) || !fs::is_directory(directory_path)) {
    std::cerr << "Error: Directory does not exist or is not a directory: "
              << directory_path << std::endl;
    return xml_paths;
  }

  // 遍历目录下的所有条目
  for (const auto &entry : fs::directory_iterator(directory_path)) {
    if (entry.is_regular_file()) {
      const fs::path &path = entry.path();

      // 检查文件扩展名是否为 .xml
      if (path.extension() == ".xml") {
        // 将完整路径添加到向量中
        xml_paths.push_back(path.stem());
      }
    }
  }

  return xml_paths;
}

int main() {
  std::vector<std::string> filename = getXmlFilePaths(xml_dir_path);
  int count = 0;
  for (auto &name : filename) {
    processImg(name);
    ++count;
    std::cerr << count << std::endl;
  }
}